# Required for Sphinx to follow autodoc_type_aliases
from __future__ import annotations

__all__ = ['Propagator', 'propagator', 'propagator_steadystate']

import numbers
import numpy as np

from .. import Qobj, qeye, qeye_like, unstack_columns, QobjEvo, liouvillian
from ..core import data as _data
from ..typing import QobjEvoLike
from .mesolve import mesolve, MESolver
from .sesolve import sesolve, SESolver
from .heom.bofin_solvers import HEOMSolver
from .solver_base import Solver
from .multitraj import MultiTrajSolver
from numbers import Number
from typing import Any
from qutip.settings import settings


def propagator_piecewise(
    H: QobjEvo,
    tlist: list[Number],
    piecewise_t: list[Number],
    c_ops: list[QobjEvo] = None,
    args: dict[str, Any] = None,
) -> Qobj | list[Qobj]:
    r"""
    Calculate propagators for piecewise constant Hamiltonians or
    Liouvillians.

    This function efficiently computes propagators when the Hamiltonian or
    Liouvillian is piecewise constant, i.e., it remains constant over
    certain time intervals and changes discontinuously at specified times.
    By exponentiating the generator directly on each constant interval,
    this approach is significantly faster than numerically integrating the
    differential equation.

    Parameters
    ----------
    H : :obj:`.QobjEvo`
        Time-dependent Hamiltonian or Liouvillian. Must be callable with
        signature ``H(t, args)`` that returns the constant operator valid at
        time ``t``.

    tlist : list of float
        List of times at which to evaluate the propagator.

    piecewise_t : list of float
        Times where the Hamiltonian or Liouvillian or collapse operators change
        discontinuously.  These times define the boundaries between constant
        intervals.
        Times outside the range ``(tlist[0], tlist[-1]]`` are ignored.

    c_ops : list of :obj:`.QobjEvo`, optional
        List of collapse operators for open quantum systems.
        Can be piecewise constant.

    args : dict, optional
        Dictionary of parameters to pass to the Hamiltonian and collapse
        operator functions.

    Returns
    -------
    propagators : list of :obj:`.Qobj`
        List of propagators :math:`U(t)` evaluated at each time in ``tlist``,
        such that :math:`\psi(t) = U(t)\psi(t_0)` for state vectors or
        :math:`\rho_{\mathrm{vec}}(t) = U(t)\rho_{\mathrm{vec}}(t_0)` for
        vectorized density matrices.
    """

    piecewise_times = {tt for tt in piecewise_t if tlist[0] < tt <= tlist[-1]}
    eval_times = set(tlist)
    times = sorted(eval_times | piecewise_times)
    atol = settings.core["atol"]

    out = []
    prev = times[0]
    prev_dt = None
    dU = None

    for nxt in times[1:]:
        dt = nxt - prev

        # Evaluate at midpoint to avoid discontinuities at boundaries
        t_eval = (prev + nxt) / 2
        H_step = H(t_eval, args)
        if c_ops:
            c_ops_q = [op(t_eval, args) for op in c_ops]
            gen = liouvillian(H_step, c_ops_q)
        else:
            gen = H_step if H_step.issuper else -1j * H_step

        if prev == times[0]:
            U = qeye_like(gen)
            if prev in eval_times:
                out.append(U)

        cannot_reuse = (
            dU is None
            or prev_dt is None
            or abs(dt - prev_dt) > atol
            or prev in piecewise_times  # Moved past a switching point
        )

        if cannot_reuse:
            dU = (gen * dt).expm()
        U = dU @ U
        prev_dt = dt

        if nxt in eval_times:
            out.append(U)
        prev = nxt
    return out


def propagator(
    H: QobjEvoLike,
    t: Number | list[Number],
    c_ops: QobjEvoLike | list[QobjEvoLike] = None,
    args: dict[str, Any] = None,
    options: dict[str, Any] = None,
    piecewise_t: list[Number] | None = None,
    **kwargs,
) -> Qobj | list[Qobj]:
    r"""
    Calculate the propagator U(t) for the density matrix or wave function such
    that :math:`\psi(t) = U(t)\psi(0)` or
    :math:`\rho_{\mathrm vec}(t) = U(t) \rho_{\mathrm vec}(0)`
    where :math:`\rho_{\mathrm vec}` is the vector representation of the
    density matrix.

    Parameters
    ----------
    H : :obj:`.Qobj`, :obj:`.QobjEvo`, :obj:`.QobjEvo` compatible format
        Possibly time-dependent system Liouvillian or Hamiltonian as a Qobj or
        QobjEvo. ``list`` of [:obj:`.Qobj`, :obj:`.Coefficient`] or callable
        that can be made into :obj:`.QobjEvo` are also accepted.

    t : float or array-like
        Time or list of times for which to evaluate the propagator. If a single
        time ``t`` is passed, the propagator from ``0`` to ``t`` is computed.
        When ``t`` is a list, the propagators from the first time in the list
        to each elements in ``t`` is returned. In that case, the first output
        will always be the identity matrix.

    c_ops : list, optional
        List of collapse operators as Qobj, QobjEvo or list that can be made
        into QobjEvo.

    args : dictionary, optional
        Parameters to callback functions for time-dependent Hamiltonians and
        collapse operators.

    options : dict, optional
        Options for the solver.

    piecewise_t : list, optional
        Times where the Hamiltonian or Liouvillian change values when
        they are piecewise constant. Providing these allows for a faster
        computation by exponentiating the Liouvillian or Hamiltonian directly
        on each interval.

    **kwargs :
        Extra parameters to use when creating the
        :obj:`.QobjEvo` from a list format ``H``. The most common are ``tlist``
        and ``order`` for array-based time dependance. See :obj:`.QobjEvo` for
        the full list.

    Returns
    -------
    U : :obj:`.Qobj`, list
        Instance representing the propagator(s) :math:`U(t)`. Return a single
        Qobj when ``t`` is a number or a list when ``t`` is a list.

    Notes
    -----
    Unlike :func:`.sesolve` or :func:`.mesolve`, the output times in ``t`` are
    not used for array time dependent system. ``tlist`` must be passed as a
    keyword argument in those case. ``tlist`` and ``t`` can have different
    length and values.

    """
    if isinstance(t, numbers.Real):
        tlist = [0, t]
        list_output = False
    else:
        tlist = t
        list_output = True

    if not isinstance(H, (Qobj, QobjEvo)):
        H = QobjEvo(H, args=args, **kwargs)

    if isinstance(c_ops, list):
        c_ops = [QobjEvo(op, args=args, **kwargs) for op in c_ops]
    elif c_ops is not None:
        c_ops = [QobjEvo(c_ops, args=args, **kwargs)]

    if piecewise_t is not None:
        out = propagator_piecewise(H, tlist, piecewise_t, c_ops, args)
        return out if list_output else out[-1]

    if c_ops:
        H = liouvillian(H, c_ops)

    U0 = qeye_like(H)

    if H.issuper:
        out = mesolve(H, U0, tlist, args=args, options=options).states
    else:
        out = sesolve(H, U0, tlist, args=args, options=options).states

    if list_output:
        return out
    else:
        return out[-1]


def propagator_steadystate(U: Qobj) -> Qobj:
    r"""Find the steady state for successive applications of the propagator
    :math:`U`.

    Parameters
    ----------
    U : :obj:`.Qobj`
        Operator representing the propagator.

    Returns
    -------
    a : :obj:`.Qobj`
        Instance representing the steady-state density matrix.
    """
    evals, estates = U.eigenstates()
    shifted_vals = np.abs(evals - 1.0)
    ev_idx = np.argmin(shifted_vals)
    rho_data = unstack_columns(estates[ev_idx].data)
    rho_data = _data.mul(rho_data, 0.5 / _data.trace(rho_data))
    return Qobj(_data.add(rho_data, _data.adjoint(rho_data)),
                dims=U._dims[0].oper,
                isherm=True,
                copy=False)


class Propagator:
    """
    A generator of propagator for a system.

    Usage:

        U = Propagator(H, c_ops)

        psi_t = U(t) @ psi_0

    Save some previously computed propagator are stored to speed up subsequent
    computation. Changing ``args`` will erase these stored probagator.

    Parameters
    ----------
    system : :obj:`.Qobj`, :obj:`.QobjEvo`, :class:`.Solver`
        Possibly time-dependent system driving the evolution, either already
        packaged in a solver, such as :class:`.SESolver` or :class:`.BRSolver`,
        or the Liouvillian or Hamiltonian as a :obj:`.Qobj`,
        :obj:`.QobjEvo`. ``list`` of [:obj:`.Qobj`, :obj:`.Coefficient`]
        or callable that can be made into :obj:`.QobjEvo` are also accepted.

        Solvers that run non-deterministacilly, such as :class:`.MCSolver`, are
        not supported.

    c_ops : list, optional
        List of :obj:`.Qobj` or :obj:`.QobjEvo` collapse operators.

    args : dictionary, optional
        Parameters to callback functions for time-dependent Hamiltonians and
        collapse operators.

    options : dict, optional
        Options for the solver.

    memoize : int, default: 10
        Max number of propagator to save.

    tol : float, default: 1e-14
        Absolute tolerance for the time. If a previous propagator was computed
        at a time within tolerance, that propagator will be returned.

    Notes
    -----
    The :class:`Propagator` is not a :obj:`.QobjEvo` so
    it cannot be used for operations with :obj:`.Qobj` or
    :obj:`.QobjEvo`. It can be made into a
    :obj:`.QobjEvo` with ::

        U = QobjEvo(Propagator(H))

    """
    def __init__(
        self,
        system: Qobj | QobjEvo | Solver,
        *,
        c_ops: QobjEvoLike | list[QobjEvoLike] = None,
        args: dict[str, Any] = None,
        options: dict[str, Any] = None,
        memoize: int = 10,
        tol: float = 1e-14,
    ):
        if isinstance(system, MultiTrajSolver):
            raise TypeError("Non-deterministic solvers cannot be used "
                            "as a propagator system")
        elif isinstance(system, HEOMSolver):
            raise NotImplementedError(
                "HEOM is not supported by Propagator. "
                "Please, tell us on GitHub issues if you need it!"
            )
        elif isinstance(system, Solver):
            self.solver = system
        else:
            Hevo = QobjEvo(system, args=args)
            c_ops = c_ops if c_ops is not None else []
            c_ops = [QobjEvo(op, args=args) for op in c_ops]
            if Hevo.issuper or c_ops:
                self.solver = MESolver(Hevo, c_ops=c_ops, options=options)
            else:
                self.solver = SESolver(Hevo, options=options)

        self.times = [0]
        self.invs = [None]
        self.props = [qeye(self.solver._sys_dims)]
        self.solver.start(self.props[0], self.times[0])
        self.cte = self.solver.rhs.isconstant
        H_0 = self.solver.rhs(0)
        self.unitary = not H_0.issuper and H_0.isherm
        self.args = args
        self.memoize = max(3, int(memoize))
        self.tol = tol

    def _lookup_or_compute(self, t):
        """
        Get U(t) from cache or compute it.
        """
        idx = np.searchsorted(self.times, t)
        if idx < len(self.times) and abs(t-self.times[idx]) <= self.tol:
            U = self.props[idx]
        elif idx > 0 and abs(t-self.times[idx-1]) <= self.tol:
            U = self.props[idx-1]
        else:
            U = self._compute(t, idx)
            self._insert(t, U, idx)
        return U

    def __call__(self, t: float, t_start: float = 0, **args):
        """
        Get the propagator from ``t_start`` to ``t``.

        Parameters
        ----------
        t : float
            Time at which to compute the propagator.
        t_start: float [0]
            Time at which the propagator start such that:
                ``psi[t] = U.prop(t, t_start) @ psi[t_start]``
        args : dict
            Argument to pass to a time dependent Hamiltonian.
            Updating ``args`` take effect since ``t=0`` and the new ``args``
            will be used in future call.
        """
        # We could improve it when the system is constant using U(2t) = U(t)**2
        if not self.cte and args and args != self.args:
            self.args = args
            self.solver._argument(args)
            self.times = [0]
            self.props = [qeye_like(self.props[0])]
            self.solver.start(self.props[0], self.times[0])

        if t_start:
            if t == t_start:
                U = self._lookup_or_compute(0)
            if self.cte:
                U = self._lookup_or_compute(t - t_start)
            else:
                Uinv = self._inv(self._lookup_or_compute(t_start))
                U = self._lookup_or_compute(t) @ Uinv
        else:
            U = self._lookup_or_compute(t)
        return U

    def inv(self, t: float, **args):
        """
        Get the inverse of the propagator at ``t``, such that
            ``psi_0 = U.inv(t) @ psi_t``

        Parameters
        ----------
        t : float
            Time at which to compute the propagator.
        args : dict
            Argument to pass to a time dependent Hamiltonian.
            Updating ``args`` take effect since ``t=0`` and the new ``args``
            will be used in future call.
        """
        return self._inv(self(t, **args))

    def _compute(self, t, idx):
        """
        Compute the propagator at ``t``, ``idx`` point to a pair of
        (time, propagator) close to the desired time.
        """
        t_last = self.solver._integrator.get_state(copy=False)[0]
        if self.times[idx-1] <= t_last <= t:
            U = self.solver.step(t)
        elif idx > 0:
            self.solver.start(self.props[idx-1], self.times[idx-1])
            U = self.solver.step(t)
        else:
            # Evolving backward in time is not supported by all integrator.
            self.solver.start(qeye_like(self.props[0]), t)
            Uinv = self.solver.step(self.times[idx])
            U = self._inv(Uinv)
        return U

    def _inv(self, U):
        return U.dag() if self.unitary else U.inv()

    def _insert(self, t, U, idx):
        """
        Insert a new pair of (time, propagator) to the memorized states.
        """
        while len(self.times) >= self.memoize:
            rm_idx = self.memoize // 2
            if self.times[rm_idx] < t:
                idx -= 1
            del self.times[rm_idx]
            del self.props[rm_idx]
        self.times.insert(idx, t)
        self.props.insert(idx, U)
